import torch
from torch import nn
from transformer_turbine_src.transformer_model import TransformerModel


class Model(nn.Module):
    def __init__(self, d_model:int, features:int, device):
        super(Model, self).__init__()
        self.transfomer = TransformerModel(features=features, d_model=d_model, device=device)
        self.output = nn.Linear(d_model, 1).to(device=device)
        self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                torch.nn.init.xavier_uniform_(m.weight)
                m.bias.data.fill_(0.01)

    def forward(self, src:torch.Tensor, tgt:torch.Tensor) -> torch.Tensor:
        y_hat = self.transfomer(src, tgt)
        y_hat = self.output(y_hat)
        return y_hat.squeeze(-1)
    
    def predict(self, src:torch.Tensor, tgt:torch.Tensor) -> torch.Tensor:
        self.eval()
        with torch.no_grad():
            y_hat = self.forward(src, tgt)
            return y_hat[:, -1]
